{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# A Practical Guide to Running CV Models: ResNet Use Case"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook serves as a practical guide to getting started running Computer Vision (CV) models on Tenstorrent hardware devices using the TT-BUDA compiler stack. *For detailed information on model compatibility, please refer to the [models support table](../model_demos/README.md#models-support-table) to check which model works with which Tenstorrent device(s).*\n",
    "\n",
    "The tutorial will walk through an example of running the [ResNet](https://en.wikipedia.org/wiki/Residual_neural_network) model on Tenstorrent AI accelerator hardware. The model weights will be directly downloaded from the [HuggingFace library](https://huggingface.co/docs/transformers/model_doc/resnet) and executed through the PyBUDA SDK.\n",
    "\n",
    "**Note on terminology:**\n",
    "\n",
    "While TT-BUDA is the official Tenstorrent AI/ML compiler stack, PyBUDA is the Python interface for TT-BUDA. TT-BUDA is the core technology; however, PyBUDA allows users to access and utilize TT-BUDA's features directly from Python. This includes directly importing model architectures and weights from PyTorch, TensorFlow, ONNX, and TFLite."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Guide Overview\n",
    "\n",
    "In this guide, we will talk through the steps for running the ResNet model trained on [ImageNet](https://www.image-net.org/) data for the **Image Classification** task.\n",
    "\n",
    "You will learn how to import the appropriate libraries, how to download model weights from popular site such as HuggingFace, utilize the PyBUDA API to initiate an inference experiment, and observe the results from running on Tenstorrent hardware."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 1: Import libraries\n",
    "\n",
    "Make sure that you have an activate Python environment with the latest version of PyBUDA installed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start by importing the pybuda library, modules from HuggingFace's transformers library, and requests, PIL, & matplotlib libraries for downloading a sample image\n",
    "import matplotlib.pyplot as plt\n",
    "import pybuda\n",
    "import requests\n",
    "from PIL import Image\n",
    "from transformers import AutoFeatureExtractor, ResNetForImageClassification"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 2: Download the model weights from HuggingFace"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load ResNet feature extractor and model from HuggingFace\n",
    "model_ckpt = \"microsoft/resnet-50\"\n",
    "feature_extractor = AutoFeatureExtractor.from_pretrained(model_ckpt)\n",
    "model = ResNetForImageClassification.from_pretrained(model_ckpt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 3: Set example input\n",
    "\n",
    "We will use a real image sample from the web. Let's stream in an image of a tiger from the ImageNet-1k dataset and view the sample."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load data sample from ImageNet-1k\n",
    "url = \"https://images.rawpixel.com/image_1300/cHJpdmF0ZS9sci9pbWFnZXMvd2Vic2l0ZS8yMDIyLTA1L3BkMTA2LTA0Ny1jaGltXzEuanBn.jpg\"\n",
    "image = Image.open(requests.get(url, stream=True).raw)\n",
    "\n",
    "# View the image\n",
    "plt.imshow(image); plt.axis(\"off\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 4: Data Preprocessing\n",
    "\n",
    "Data preprocessing is an important step in the AI inference pipeline. For CV models, we apply transformations to the input image such as centering, cropping, padding, resizing, scaling, and normalizing. Some libraries, such as HuggingFace's transformers and PyTorch Image Models (timm), have transform classes to handle this for you."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Data preprocessing\n",
    "pixel_values = feature_extractor(image, return_tensors=\"pt\")[\"pixel_values\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 5: Configure PyBUDA Parameters\n",
    "\n",
    "There are optional configurations that can be adjusted before compiling and running a model on Tenstorrent hardware. Sometimes, the configurations are necessary to compile the model and other times they are tuneable parameters that can be adjusted for performance enhancement.\n",
    "\n",
    "For the ResNet model, two key parameters are required for compilation:\n",
    "\n",
    "* `balancer_policy`\n",
    "* `enable_t_streaming`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set PyBUDA configuration parameters\n",
    "compiler_cfg = pybuda.config._get_global_compiler_config()  # get global configuration object\n",
    "compiler_cfg.balancer_policy = \"Ribbon\"  # set balancer policy\n",
    "compiler_cfg.enable_t_streaming = True  # enable tensor streaming"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 6: Instantiate Tenstorrent device\n",
    "\n",
    "The first time we use PyBUDA, we must initialize a `TTDevice` object which serves as the abstraction over the target hardware."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tt0 = pybuda.TTDevice(\n",
    "    name=\"tt_device_0\",  # here we can give our device any name we wish, for tracking purposes\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 7: Create a PyBUDA module from PyTorch model\n",
    "\n",
    "Next, we must abstract the PyTorch model loaded from HuggingFace into a `pybuda.PyTorchModule` object. This will let the BUDA compiler know which model architecture and AI framework it has to compile.\n",
    "\n",
    "We then \"place\" this module onto the previously initialized `TTDevice`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create module\n",
    "pybuda_module = pybuda.PyTorchModule(\n",
    "    name = \"pt_resnet50\",  # give the module a name, this will be used for tracking purposes\n",
    "    module=model  # specify the model that is being targeted for compilation\n",
    ")\n",
    "\n",
    "# Place module on device\n",
    "tt0.place_module(module=pybuda_module)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 8: Push the inputs into the model input queue"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Push inputs\n",
    "tt0.push_to_inputs((pixel_values,))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 9: Run inference on the targeted device\n",
    "\n",
    "Running a model on a Tenstorrent device invovles two parts: compilation and runtime.\n",
    "\n",
    "Compilation -- TT-BUDA is a compiler. Meaning that it will take a model architecture graph and compile it for the target hardware. Compilation can take anywhere from a few seconds to a few minutes, depending on the model. This only needs to happen once. When you execute the following block of code the compilation logs will be displayed.\n",
    "\n",
    "Runtime -- once the model has been compiled and loaded onto the device, the user can push new inputs which will execute immediately.\n",
    "\n",
    "The `run_inference` API can achieve both steps in a single call. If it's the first call, the model will compile. Any subsequent calls will execute runtime only.\n",
    "\n",
    "Please refer to the documentation for alternative APIs such as `initialize_pipeline` and `run_forward`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run inference on Tenstorrent device\n",
    "output_q = pybuda.run_inference()  # executes compilation (if first time) + runtime\n",
    "output = output_q.get()  # get last value from output queue"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 10: Data Postprocessing\n",
    "\n",
    "Data postprocessing is done to convert the model outputs into a readable / useful format. For image classification tasks, this usually means receiving the logit outputs from the model, extracting the top predicted class, and matching this with an entry from the label dictionary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Data postprocessing\n",
    "predicted_value = output[0].value().argmax(-1).item()\n",
    "predicted_label = model.config.id2label[predicted_value]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 11: Print and evaluate outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Print outputs\n",
    "print(f\"Predicted_label: {predicted_label}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 12: Shutdown PyBuda"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pybuda.shutdown()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
